# Copyright (c) 2024 Huawei Technologies Co., Ltd.
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This code is inspired by the pytorch's torchtune library.
# https://github.com/pytorch/torchtune
#
# openMind is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#
#          http://license.coscl.org.cn/MulanPSL2
#
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.
import argparse
import os
import re
import string
from urllib.parse import urlparse
from typing import List, Union

import yaml
from tabulate import tabulate

from openmind.utils import is_vision_available, logging

if is_vision_available():
    from PIL.Image import Image

logger = logging.get_logger(__name__)


def _trans_args_list_to_dict(args_list: List) -> dict:

    if len(args_list) % 2 != 0:
        raise ValueError(f"arguments in CLI should be specified in paris, but got {args_list}")

    args_dict = {}

    for idx in range(0, len(args_list), 2):
        key = args_list[idx]
        val = args_list[idx + 1]

        if not key.startswith("--"):
            raise ValueError(f"Argument name `{key}` in CLI should start with `--`")

        if val.startswith("--"):
            raise ValueError(f"Argument value `{val}` of key `{key}` in CLI should not start with `--`")

        args_dict[key.lstrip("-")] = val

    return args_dict


def str2bool(value):

    if isinstance(value, bool):
        return value
    if value.lower() in {"true", "1"}:
        return True
    elif value.lower() in {"false", "0"}:
        return False
    else:
        raise argparse.ArgumentTypeError("Boolean value expected.")


def str2bool_or_auto(value):
    if isinstance(value, bool):
        return value
    value = value.lower()
    if value in {"true", "yes", "1"}:
        return True
    elif value in {"false", "no", "0"}:
        return False
    elif value == "auto":
        return "auto"
    else:
        raise argparse.ArgumentTypeError("Value should be one of: 'true', 'false', or 'auto'.")


def looks_like_url(s: str) -> bool:
    """check url format"""
    _SCHEME_RE = re.compile(r"^[a-zA-Z][a-zA-Z0-9+.-]*$")  # RFC3986 scheme
    if not isinstance(s, str) or not s:
        return False
    try:
        p = urlparse(s)
    except Exception:
        return False
    return bool(p.scheme and p.netloc and _SCHEME_RE.match(p.scheme))


def validate_directory(path, allow_symlinks=True):
    """
    Validate directory path with realpath resolution
    """
    if not isinstance(path, str) or not path.strip():
        raise ValueError("Directory path must be a non-empty string")

    # Normalize and resolve real path
    normalized_path = os.path.normpath(path.strip())
    real_path = os.path.realpath(normalized_path)

    # Symlink policy
    if not allow_symlinks:
        if os.path.islink(normalized_path) or real_path != normalized_path:
            raise ValueError(f"Symbolic link not allowed: {path} -> {real_path}")
    elif os.path.islink(normalized_path):
        logger.warning(f"Directory path is a symbolic link: {normalized_path} -> {real_path}")

    # Check if the resolved path is a directory
    if not os.path.isdir(real_path):
        raise ValueError(f"Path is not a directory or not exists: {real_path}")

    # Character validation on the resolved path

    # Component pattern check
    pattern = r"^(/|\./?|[~0-9a-zA-Z_\-\s\./\u4e00-\u9fa5])+$"
    if not re.fullmatch(pattern, real_path):
        raise RuntimeError(f"Invalid directory path: {real_path}")

    return real_path


def validate_cache_dir(cache_dir):
    """A function specifically for validating the cache directory, allowing non-existent paths"""
    if not isinstance(cache_dir, str):
        raise ValueError(f"Cache directory must be a string, got {type(cache_dir)}")

    # for cache, we allow non-existent paths
    normalized_path = os.path.normpath(cache_dir.strip())

    # basic check
    pattern = r"^(/|\./?|[~0-9a-zA-Z_\-\s\./\u4e00-\u9fa5])+$"

    if not re.fullmatch(pattern, normalized_path):
        raise ValueError(f"Invalid cache directory path format: {cache_dir}")

    return normalized_path


def _validate_file_path_common(
    file_path: str,
    allowed_extensions: list = None,
    allow_symlinks=True,
    require_read=True,
    require_write=True,
):
    """
    Universal file path validation function

    Args:
        file_path: File path to validate
        allowed_extensions: List of allowed file extensions
        allow_symlinks: Whether to allow symbolic links
        require_read: Whether to require read access to the file
        require_write: Whether to require write access to the file
        component_pattern: Regular expression pattern to validate file path components
    Returns:
        str: Normalized file path
    """
    normalized_path = os.path.normpath(file_path)
    real_path = os.path.realpath(normalized_path)

    # Symlink policy: if disallowed, reject when resolved path differs or original (or any component) is a symlink
    if not allow_symlinks:
        if os.path.islink(normalized_path) or real_path != normalized_path:
            raise ValueError(f"Symbolic link not allowed: {file_path} -> {real_path}")

    # Must be an existing regular file
    if not os.path.isfile(real_path):
        # isfile implicitly checks existence
        raise FileNotFoundError(f"File not found or not a regular file: {file_path}")

    # Extension whitelist
    if allowed_extensions:
        normalized_extensions = [ext if ext.startswith(".") else f".{ext}" for ext in allowed_extensions]
        if not any(real_path.lower().endswith(ext.lower()) for ext in normalized_extensions):
            raise ValueError(f"File must have one of extensions: {normalized_extensions}")

    # Permissions
    if require_read and not os.access(real_path, os.R_OK):
        raise PermissionError(f"No read permission: {real_path}")
    if require_write and not os.access(real_path, os.W_OK):
        raise PermissionError(f"No write permission: {real_path}")

    pattern = r"^(/|\./?|[~0-9a-zA-Z_\-\s\./\u4e00-\u9fa5])+$"
    if not re.fullmatch(pattern, real_path):
        raise RuntimeError(f"Invalid input path: {real_path}")
    return real_path


def safe_load_yaml(path):
    if path is None:
        raise ValueError("param `path` is required for [safe_load_yaml]")

    if not isinstance(path, str):
        raise TypeError(f"param `path` should be string format for `safe_load_yaml func`, but got {type(path)} type.")

    normalized_path = _validate_file_path_common(path, allowed_extensions=[".yaml", ".yml"], require_read=True)

    with open(normalized_path, "r") as file:
        content = yaml.safe_load(file)

    return content


def validate_file_path(
    file_path,
    allowed_extensions: List = None,
    allow_symlinks: bool = True,
    require_read: bool = True,
    require_write: bool = True,
):
    """
    Validate the security of a file path

    Args:
        file_path: File path to validate
        allowed_extensions: List of allowed file extensions
        allow_symlinks: Whether to allow symbolic links
    """
    # Directly call the common function
    return _validate_file_path_common(file_path, allowed_extensions, allow_symlinks, require_read, require_write)


def validate_image_path(image_path):
    """Validate the security of image paths"""
    if not isinstance(image_path, str):
        # If it's not a string (e.g. already a PIL image object), return directly
        return image_path

    # If it looks like a URL, perform URL validation
    if looks_like_url(image_path):
        return validate_url(image_path)

    # Otherwise, treat it as a file path and validate it
    # Common image file extensions supported by PIL/Pillow
    image_extensions = [
        ".png",
        ".jpg",
        ".jpeg",
        ".gif",
        ".bmp",
        ".tiff",
        ".tif",
        ".webp",
        ".ppm",
        ".pgm",
        ".pbm",
        ".pnm",
        ".ico",
        ".icns",
        ".eps",
        ".psd",
        ".xcf",
        ".pcx",
        ".sgi",
        ".rgb",
        ".rgba",
        ".j2k",
        ".jp2",
        ".jpx",
        ".jpf",
        ".jpm",
        ".mj2",
    ]
    try:
        return _validate_file_path_common(image_path, allowed_extensions=image_extensions, require_read=True)
    except FileNotFoundError as e:
        # Re-raise with a more specific error message
        raise FileNotFoundError(
            f"Image path '{image_path}' is not a valid URL and was not found as a local file."
        ) from e


def validate_image_path_list(image: Union[str, List[str], "Image", List["Image"]]):
    # Perform security validation on image input
    if isinstance(image, str):
        image = validate_image_path(image)
    elif isinstance(image, List):
        # If it's a list, validate each element
        validated_images = []
        for img in image:
            img = validate_image_path(img)
            validated_images.append(img)
        image = validated_images
    return image


def validate_url(url: str) -> str:
    """
    Perform only character-level legality checks for a URL-like string.
    This does NOT validate structure (netloc, path semantics, reachability).
    Checks:
      - non-empty str
      - no control / whitespace chars
      - scheme (if present) matches ^[A-Za-z][A-Za-z0-9+.-]*$
      - percent-encoding tokens are well-formed (every '%' followed by two hex digits)
      - all chars belong to a conservative allowed set (RFC 3986 unreserved + reserved)
    Returns the original string if all checks pass; raises ValueError otherwise.
    """
    if not isinstance(url, str) or not url:
        raise ValueError("Invalid URL: must be a non-empty string")

    # Reject control chars and spaces
    if re.search(r"[\x00-\x20\x7F]", url):
        raise ValueError(f"Invalid URL: contains control or whitespace characters: {url!r}")

    # Percent-encoding correctness
    if re.search(r"%(?![0-9A-Fa-f]{2})", url):
        raise ValueError(f"Invalid URL: malformed percent-encoding: {url!r}")

    # Optional scheme check (only if a scheme appears)
    if ":" in url:
        _SCHEME_CHARS = set(string.ascii_letters + string.digits + "+-.")

        def _valid_scheme(s: str) -> bool:
            return bool(s) and s[0] in string.ascii_letters and all(c in _SCHEME_CHARS for c in s)

        head, sep, _ = url.partition(":")
        if sep:
            scheme = head
            if scheme and not _valid_scheme(scheme):
                raise ValueError(f"Invalid URL: invalid scheme: {scheme!r}")

    if url.endswith(":") or url.endswith("://"):
        raise ValueError(f"Invalid URL: incomplete URL structure: {url!r}")

    # Allowed characters (unreserved + reserved + percent-escapes + '#')
    # unreserved: A-Z a-z 0-9 - . _ ~
    # reserved: : / ? # [ ] @ ! $ & ' ( ) * + , ; =
    allowed_pattern = r"^[A-Za-z0-9\-._~:/?#\[\]@!$&\'()*+,;=%]*$"
    if not re.fullmatch(allowed_pattern, url):
        raise ValueError(f"Invalid URL: contains disallowed characters: {url!r}")

    return url


def validate_input_or_path(input_or_path):
    """
    Validate input parameter, which can be a URL, local file path, or string

    Args:
        input_or_path: Input parameter, could be a URL, file path, or plain string

    Returns:
        Validated input parameter

    Raises:
        ValueError: Raised when the input parameter is invalid
    """
    if not isinstance(input_or_path, (str, dict)):
        raise ValueError(f"Input must be a string or dict, got {type(input_or_path)}")

    # If it's a dict, return directly
    if isinstance(input_or_path, dict):
        return input_or_path

    # If it's a string, perform further validation
    if isinstance(input_or_path, str):
        # First check if it's a valid URL
        try:
            validated_url = validate_url(input_or_path)
            return validated_url
        except ValueError:
            # Not a valid URL, continue to check if it's a local file
            pass

        # Check if it's a local file or directory path
        if os.path.exists(input_or_path):
            # If it's a directory, validate it as a directory
            if os.path.isdir(input_or_path):
                return validate_directory(input_or_path)
            # If it's a file, validate it as a file
            else:
                return validate_file_path(input_or_path)

        # If it's neither a URL nor an existing file, treat it as a plain string input (like a prompt)
        # In this case, no special validation is needed, return directly
        return input_or_path

    return input_or_path


def print_formatted_table(data, header, missingval="N/A"):
    print(tabulate(data, header, missingval=missingval, tablefmt="fancy_grid"))
